package com.cloud.gateway.filter;

import com.cloud.common.loadbalance.DefaultLoadBalancer;
import com.cloud.common.loadbalance.ServiceInstanceStatus;
import com.cloud.common.loadbalance.ServiceInstanceStatusManager;
import lombok.extern.slf4j.Slf4j;
import org.springframework.cloud.client.ServiceInstance;
import org.springframework.cloud.client.loadbalancer.LoadBalancerProperties;
import org.springframework.cloud.client.loadbalancer.LoadBalancerUriTools;
import org.springframework.cloud.gateway.config.GatewayLoadBalancerProperties;
import org.springframework.cloud.gateway.filter.GatewayFilterChain;
import org.springframework.cloud.gateway.filter.ReactiveLoadBalancerClientFilter;
import org.springframework.cloud.gateway.support.DelegatingServiceInstance;
import org.springframework.cloud.gateway.support.NotFoundException;
import org.springframework.cloud.loadbalancer.support.LoadBalancerClientFactory;
import org.springframework.web.server.ServerWebExchange;
import reactor.core.publisher.Mono;

import java.net.URI;

import static org.springframework.cloud.gateway.support.ServerWebExchangeUtils.*;

/**
 * 均衡负载过滤器
 *
 * @author fengwenjin
 */
@Slf4j
public class LoadBalancerFilter extends ReactiveLoadBalancerClientFilter {
    public static final String SERVICE_INSTANCE = "ServiceInstance";

    public LoadBalancerFilter(LoadBalancerClientFactory clientFactory, GatewayLoadBalancerProperties properties, LoadBalancerProperties loadBalancerProperties) {
        super(clientFactory, properties, loadBalancerProperties);
    }

    @Override
    public Mono<Void> filter(ServerWebExchange exchange, GatewayFilterChain chain) {
        URI url = exchange.getAttribute(GATEWAY_REQUEST_URL_ATTR);
        String schemePrefix = exchange.getAttribute(GATEWAY_SCHEME_PREFIX_ATTR);
        if (url == null || (!"lb".equals(url.getScheme()) && !"lb".equals(schemePrefix))) {
            return chain.filter(exchange);
        }
        // preserve the original url
        addOriginalRequestUrl(exchange, url);

        URI uri = exchange.getAttribute(GATEWAY_REQUEST_URL_ATTR);
        ServiceInstance instance = DefaultLoadBalancer.choose(uri.getHost());
        if (instance == null) {
            String message = "Unable to find instance for " + url.getHost();
            throw new NotFoundException(message);
        }

        String overrideScheme = instance.isSecure() ? "https" : "http";
        if (schemePrefix != null) {
            overrideScheme = url.getScheme();
        }

        DelegatingServiceInstance serviceInstance = new DelegatingServiceInstance(instance, overrideScheme);
        URI requestUrl = LoadBalancerUriTools.reconstructURI(serviceInstance, uri);

        if (log.isTraceEnabled()) {
            log.trace("LoadBalancerFilter url chosen: " + requestUrl);
        }
        // 统计实例状态
        ServiceInstanceStatus instanceStatus = ServiceInstanceStatusManager.get(serviceInstance);
        instanceStatus.incrementActiveRequestsCount();

        exchange.getAttributes().put(SERVICE_INSTANCE, serviceInstance);
        exchange.getAttributes().put(GATEWAY_REQUEST_URL_ATTR, requestUrl);
        return chain.filter(exchange);
    }
}
